package com.zfc.threadpool.config.thread;

import com.zfc.threadpool.enums.RejectedPolicyType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableAsync;

import java.util.concurrent.*;

/**
 * @author zfc
 * @version 1.0
 */
@Slf4j
@EnableAsync
@Configuration
@EnableConfigurationProperties(ThreadPoolConfigProperties.class)
public class ThreadPoolConfig {

    /**
     *  Comfyui计算所需线程池
     * @param properties 线程池配置类
     * @return 线程池Bean
     */
    @Bean("comfyUIThreadExecutor")
    public ThreadPoolExecutor comfyUIThreadExecutor(ThreadPoolConfigProperties properties) {
        return new ThreadPoolExecutor(
            properties.getCorePoolSize(),
            properties.getMaxPoolSize(),
            properties.getKeepAliveTime(),
            TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(properties.getBlockQueueSize()),
            Executors.defaultThreadFactory(),
            createRejectedPolicy(properties.getPolicy())
        );
    }

    /**
     *  拒绝策略
     * @param policyType 拒绝策略类型
     * @return 拒绝策略
     */
    private RejectedExecutionHandler createRejectedPolicy(String policyType) {
        return RejectedPolicyType.of(policyType).getHandler();
    }
}
